#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
可视化标注图片生成工具
根据 images 目录中的文件和 annotations 中的标注内容生成带标注框的可视化图片

执行指令：
conda activate seal_env && \
python continue_gen_va_pic.py \
    --dataset_path /path/to/dataset \
    --output_dir va

功能说明:
    - 读取 images 目录中的所有图片文件
    - 读取 annotations/instance_train.json 和 annotations/instance_val.json 中的标注信息
    - 根据标注信息在图片上绘制边界框和标签
    - 生成的可视化图片保存到 va 目录，文件名添加 _res 后缀
    - 支持 COCO 格式的标注数据

使用方法:
    # 使用默认参数
    python continue_gen_va_pic.py --dataset_path dataset
    
    # 指定输出目录
    python continue_gen_va_pic.py --dataset_path dataset --output_dir visualized_annotations
    
    # 查看帮助信息
    python continue_gen_va_pic.py --help

参数说明:
    --dataset_path: 数据集路径，包含 images 和 annotations 目录
    --output_dir: 可视化图片输出目录，默认为 va

输出格式:
    生成的可视化图片包含:
    - 原始图片作为背景
    - 检测框（矩形边界框）
    - 类别标签和置信度
    - 不同类别使用不同颜色

注意事项:
    - 需要安装PIL库: pip install Pillow
    - 支持常见图片格式：PNG、JPEG、BMP、TIFF等
    - 可视化图片格式为JPG
"""

import os
import json
import argparse
import glob
from typing import Dict, List, Any, Tuple
from pathlib import Path

# 尝试导入图像处理库
PIL_AVAILABLE = False
OPENCV_AVAILABLE = False

try:
    from PIL import Image, ImageDraw, ImageFont
    PIL_AVAILABLE = True
    print("✓ PIL/Pillow库可用")
except ImportError:
    try:
        import cv2
        OPENCV_AVAILABLE = True
        print("✓ OpenCV库可用")
    except ImportError:
        PIL_AVAILABLE = False
        OPENCV_AVAILABLE = False
        print("⚠ 警告: 图像处理库未安装，可视化功能将不可用")
        print("请运行: pip install Pillow 或 pip install opencv-python")

class VisualizationGenerator:
    """可视化标注图片生成器"""
    
    def __init__(self):
        # COCO类别定义
        self.categories = {
            0: {"name": "_background_", "color": (128, 128, 128)},
            1: {"name": "Text", "color": (255, 0, 0)},           # 红色
            2: {"name": "Title", "color": (0, 255, 0)},          # 绿色
            3: {"name": "Figure", "color": (0, 0, 255)},        # 蓝色
            4: {"name": "Figure caption", "color": (255, 255, 0)}, # 青色
            5: {"name": "Table", "color": (255, 0, 255)},       # 洋红
            6: {"name": "Table caption", "color": (0, 255, 255)}, # 黄色
            7: {"name": "Header", "color": (128, 0, 128)},       # 紫色
            8: {"name": "Footer", "color": (128, 128, 0)},       # 橄榄色
            9: {"name": "Seal", "color": (255, 165, 0)},        # 橙色
            10: {"name": "Equation", "color": (0, 128, 128)},    # 深青色
            11: {"name": "Handwriting", "color": (255, 192, 203)} # 粉色
        }
        
        # 字体设置
        self.font_size = 12
        self.font = None
        if PIL_AVAILABLE:
            try:
                # 尝试使用系统字体
                self.font = ImageFont.truetype("/System/Library/Fonts/Arial.ttf", self.font_size)
            except:
                try:
                    # 尝试使用默认字体
                    self.font = ImageFont.load_default()
                except:
                    self.font = None
    
    def load_annotations(self, dataset_path: str) -> Dict[str, Any]:
        """加载标注数据"""
        print("正在加载标注数据...")
        
        all_images = []
        all_annotations = []
        
        # 加载训练数据
        train_file = os.path.join(dataset_path, "annotations", "instance_train.json")
        if os.path.exists(train_file):
            with open(train_file, 'r', encoding='utf-8') as f:
                train_data = json.load(f)
                all_images.extend(train_data.get('images', []))
                all_annotations.extend(train_data.get('annotations', []))
            print(f"✓ 加载训练数据: {len(train_data.get('images', []))} 个图片")
        
        # 加载验证数据
        val_file = os.path.join(dataset_path, "annotations", "instance_val.json")
        if os.path.exists(val_file):
            with open(val_file, 'r', encoding='utf-8') as f:
                val_data = json.load(f)
                all_images.extend(val_data.get('images', []))
                all_annotations.extend(val_data.get('annotations', []))
            print(f"✓ 加载验证数据: {len(val_data.get('images', []))} 个图片")
        
        # 创建图片ID到图片信息的映射
        image_id_to_info = {img['id']: img for img in all_images}
        
        # 按图片分组标注
        annotations_by_image = {}
        for annotation in all_annotations:
            image_id = annotation['image_id']
            if image_id not in annotations_by_image:
                annotations_by_image[image_id] = []
            annotations_by_image[image_id].append(annotation)
        
        print(f"✓ 总计: {len(all_images)} 个图片, {len(all_annotations)} 个标注")
        
        return {
            'images': all_images,
            'annotations_by_image': annotations_by_image,
            'image_id_to_info': image_id_to_info
        }
    
    def get_image_files(self, images_dir: str) -> List[str]:
        """获取所有图片文件"""
        image_extensions = ['*.png', '*.jpg', '*.jpeg', '*.bmp', '*.tiff', '*.tif', '*.gif', '*.webp']
        image_files = []
        
        for ext in image_extensions:
            pattern = os.path.join(images_dir, ext)
            image_files.extend(glob.glob(pattern))
        
        return sorted(image_files)
    
    def draw_annotations_pil(self, image, annotations: List[Dict], image_info: Dict):
        """使用PIL绘制标注"""
        draw = ImageDraw.Draw(image)
        
        for annotation in annotations:
            # 获取边界框信息
            bbox = annotation['bbox']  # [x, y, width, height]
            category_id = annotation['category_id']
            
            # 获取类别信息
            if category_id in self.categories:
                category_name = self.categories[category_id]['name']
                color = self.categories[category_id]['color']
            else:
                category_name = f"Unknown_{category_id}"
                color = (128, 128, 128)
            
            # 绘制边界框
            x, y, w, h = bbox
            x1, y1 = int(x), int(y)
            x2, y2 = int(x + w), int(y + h)
            
            # 绘制矩形框
            draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
            
            # 绘制标签
            label = f"{category_name}"
            if self.font:
                # 获取文本尺寸
                bbox_text = draw.textbbox((0, 0), label, font=self.font)
                text_width = bbox_text[2] - bbox_text[0]
                text_height = bbox_text[3] - bbox_text[1]
            else:
                text_width = len(label) * 6
                text_height = 12
            
            # 绘制标签背景
            label_y = max(y1 - text_height - 2, 0)
            draw.rectangle([x1, label_y, x1 + text_width + 4, label_y + text_height + 2], 
                          fill=color, outline=color)
            
            # 绘制标签文字
            text_color = (255, 255, 255) if sum(color) < 400 else (0, 0, 0)
            draw.text((x1 + 2, label_y + 1), label, fill=text_color, font=self.font)
        
        return image
    
    def draw_annotations_opencv(self, image, annotations: List[Dict], image_info: Dict):
        """使用OpenCV绘制标注"""
        for annotation in annotations:
            # 获取边界框信息
            bbox = annotation['bbox']  # [x, y, width, height]
            category_id = annotation['category_id']
            
            # 获取类别信息
            if category_id in self.categories:
                category_name = self.categories[category_id]['name']
                color = self.categories[category_id]['color']
            else:
                category_name = f"Unknown_{category_id}"
                color = (128, 128, 128)
            
            # 绘制边界框
            x, y, w, h = bbox
            x1, y1 = int(x), int(y)
            x2, y2 = int(x + w), int(y + h)
            
            # 绘制矩形框
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            
            # 绘制标签
            label = f"{category_name}"
            font_scale = 0.5
            thickness = 1
            
            # 获取文本尺寸
            (text_width, text_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
            
            # 绘制标签背景
            label_y = max(y1 - text_height - 2, text_height)
            cv2.rectangle(image, (x1, label_y - text_height - 2), (x1 + text_width + 4, label_y + 2), color, -1)
            
            # 绘制标签文字
            cv2.putText(image, label, (x1 + 2, label_y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness)
        
        return image
    
    def process_single_image(self, image_path: str, annotations: List[Dict], image_info: Dict, output_path: str):
        """处理单个图片"""
        try:
            # 检查是否有图像处理库
            if not PIL_AVAILABLE and not OPENCV_AVAILABLE:
                print(f"⚠ 跳过图片处理（无图像处理库）: {os.path.basename(image_path)}")
                # 创建一个简单的文本文件作为占位符
                with open(output_path.replace('.jpg', '.txt'), 'w') as f:
                    f.write(f"可视化图片: {os.path.basename(image_path)}\n")
                    f.write(f"标注数量: {len(annotations)}\n")
                    for i, ann in enumerate(annotations):
                        f.write(f"标注{i+1}: 类别{ann['category_id']}, 边界框{ann['bbox']}\n")
                return True
            
            # 读取图片
            if PIL_AVAILABLE:
                image = Image.open(image_path)
                if image.mode != 'RGB':
                    image = image.convert('RGB')
                
                # 绘制标注
                image_with_annotations = self.draw_annotations_pil(image, annotations, image_info)
                
                # 保存图片
                image_with_annotations.save(output_path, 'JPEG', quality=95)
                
            elif OPENCV_AVAILABLE:
                image = cv2.imread(image_path)
                if image is None:
                    print(f"✗ 无法读取图片: {image_path}")
                    return False
                
                # 绘制标注
                image_with_annotations = self.draw_annotations_opencv(image, annotations, image_info)
                
                # 保存图片
                cv2.imwrite(output_path, image_with_annotations)
            
            print(f"✓ 生成可视化图片: {os.path.basename(output_path)}")
            return True
            
        except Exception as e:
            print(f"✗ 处理图片失败 {image_path}: {e}")
            return False
    
    def generate_visualizations(self, dataset_path: str, output_dir: str):
        """生成所有可视化图片"""
        print("=" * 60)
        print("开始生成可视化标注图片")
        print("=" * 60)
        
        # 检查输入目录
        images_dir = os.path.join(dataset_path, "images")
        if not os.path.exists(images_dir):
            print(f"✗ 错误: images目录不存在: {images_dir}")
            return False
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        print(f"✓ 输出目录: {output_dir}")
        
        # 加载标注数据
        annotation_data = self.load_annotations(dataset_path)
        if not annotation_data['images']:
            print("✗ 错误: 没有找到标注数据")
            return False
        
        # 获取所有图片文件
        image_files = self.get_image_files(images_dir)
        print(f"✓ 找到 {len(image_files)} 个图片文件")
        
        if not image_files:
            print("✗ 错误: 没有找到图片文件")
            return False
        
        # 创建文件名到图片信息的映射
        filename_to_info = {}
        for img_info in annotation_data['images']:
            filename_to_info[img_info['file_name']] = img_info
        
        processed_count = 0
        skipped_count = 0
        
        # 处理每个图片文件
        for image_path in image_files:
            filename = os.path.basename(image_path)
            
            # 查找对应的图片信息
            if filename in filename_to_info:
                image_info = filename_to_info[filename]
                image_id = image_info['id']
                
                # 获取该图片的标注
                annotations = annotation_data['annotations_by_image'].get(image_id, [])
                
                if annotations:
                    # 生成输出文件名
                    name, ext = os.path.splitext(filename)
                    output_filename = f"{name}_res.jpg"
                    output_path = os.path.join(output_dir, output_filename)
                    
                    # 处理图片
                    if self.process_single_image(image_path, annotations, image_info, output_path):
                        processed_count += 1
                    else:
                        skipped_count += 1
                else:
                    print(f"⚠ 跳过无标注的图片: {filename}")
                    skipped_count += 1
            else:
                print(f"⚠ 跳过未找到标注信息的图片: {filename}")
                skipped_count += 1
        
        print("\n" + "=" * 60)
        print("可视化图片生成完成！")
        print(f"✓ 成功处理: {processed_count} 个图片")
        print(f"⚠ 跳过: {skipped_count} 个图片")
        print(f"✓ 输出目录: {output_dir}")
        print("=" * 60)
        
        return True


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='可视化标注图片生成工具')
    parser.add_argument('--dataset_path', 
                       required=True,
                       help='数据集路径，包含 images 和 annotations 目录')
    parser.add_argument('--output_dir', 
                       default='va',
                       help='可视化图片输出目录 (默认: va)')
    
    args = parser.parse_args()
    
    # 检查数据集路径是否存在
    if not os.path.exists(args.dataset_path):
        print(f"错误: 数据集路径不存在: {args.dataset_path}")
        return
    
    # 检查必要的目录
    images_dir = os.path.join(args.dataset_path, "images")
    annotations_dir = os.path.join(args.dataset_path, "annotations")
    
    if not os.path.exists(images_dir):
        print(f"错误: images目录不存在: {images_dir}")
        return
    
    if not os.path.exists(annotations_dir):
        print(f"错误: annotations目录不存在: {annotations_dir}")
        return
    
    if args.output_dir == 'va':
        args.output_dir = os.path.join(args.dataset_path, "va")
    
    # 创建生成器并处理数据
    generator = VisualizationGenerator()
    success = generator.generate_visualizations(args.dataset_path, args.output_dir)
    
    if success:
        print("\n可视化图片生成成功！")
    else:
        print("\n可视化图片生成失败！")
        exit(1)


if __name__ == "__main__":
    main()
